package rddl.policy.expectation;

import org.apache.commons.math3.random.RandomDataGenerator;
import rddl.ActionGenerator;
import rddl.EvalException;
import rddl.RDDL;
import rddl.State;
import rddl.policy.Policy;
import util.Pair;

import java.util.*;

/**
 * Created by HuGuodong on 2017/11/24.
 */

public class MaxExpectationPolicy extends Policy {

    public final static boolean DISPLAY_UPDATES = false;

    public Random _rand = new Random();

    // Local vars
    public RDDL.INSTANCE _rddlInstance = null;
    public int _nHorizon;
    public ArrayList<RDDL.PVAR_NAME> _alStateNames = new ArrayList<RDDL.PVAR_NAME>();
    public ArrayList<RDDL.PVAR_NAME> _alActionNames = new ArrayList<RDDL.PVAR_NAME>();
    public TreeMap<Pair, RDDL.PVAR_NAME> _tmIntermNames = new TreeMap<Pair, RDDL.PVAR_NAME>();
    public ArrayList<RDDL.PVAR_NAME> _alIntermNames = new ArrayList<RDDL.PVAR_NAME>();
    public ArrayList<RDDL.PVAR_NAME> _alObservNames = new ArrayList<RDDL.PVAR_NAME>();
    public ArrayList<RDDL.PVAR_NAME> _alNonFluentNames = new ArrayList<RDDL.PVAR_NAME>();
    public HashMap<String,ArrayList<RDDL.PVAR_NAME>> _hmTypeMap = new HashMap<String,ArrayList<RDDL.PVAR_NAME>>();

    // PVariable definitions
    public HashMap<RDDL.PVAR_NAME, RDDL.PVARIABLE_DEF> _hmPVariables;

    // CPF definitions
    public HashMap<RDDL.PVAR_NAME, RDDL.CPF_DEF> _hmCPFs;

    public MaxExpectationPolicy() { }

    public MaxExpectationPolicy(String instance_name) {
        super(instance_name);
    }

    public void init(){
    }

    /**
     * override getActions
     * 1. 暂时只计算下一个状态动作reward的期望
     * @param s
     * @return
     * @throws EvalException
     */
    public ArrayList<RDDL.PVAR_INST_DEF> getActions(State s) throws EvalException{
        init();
        if (s == null) {
            return new ArrayList<RDDL.PVAR_INST_DEF>();
        }

        // Get a map of { legal action names -> RDDL action definition }
        Map<String,ArrayList<RDDL.PVAR_INST_DEF>> action_map =
                ActionGenerator.getLegalBoolActionMap(s);

        // Return a random action selection
        ArrayList<String> actions = new ArrayList<String>(action_map.keySet());

        // save action : reward expectation pair
        Map<String,Double> _mActionRewardException = new HashMap<>();

        // calculate each action's reward expectation
        for (String action_simulate : actions) {

            // clone a current state, which will be used for further calculation.
            State mepState = Util.coloneAndExtendState(s);

            RandomDataGenerator _randGen = new RandomDataGenerator();

            Double rewardExpectation = mepState.computeNextStateReward(action_map.get(action_simulate), _randGen);

            _mActionRewardException.put(action_simulate, rewardExpectation);
        }

        System.out.println("action expectation "+_mActionRewardException);

        // find maxExpectationAction: return action:reward expectation pair
        Map<String,Double> maxActionExpectation = Util.maxExpectationAction(_mActionRewardException);

        //  action_token = maxActionExpectation
        String action_taken = null;
        for (String action : maxActionExpectation.keySet()) {
            System.out.println("action choose: action = " + action +
                                " and reward expectation = " + maxActionExpectation.get(action));
            action_taken = action;
        }

        System.out.println("\n--> Action taken: " + action_taken);

        return action_map.get(action_taken);
    }

}
